[TRTLLM-14715][feat] preserve native MoE A2A graph VAs across restore - #16632

tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py (1)

703-722: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider mirroring the _require_mapped() guard here for cross-implementation consistency.

MoeAlltoAll.get_combine_payload_tensor_in_workspace() calls self._require_mapped() before the phase check, but this NVLink one-sided variant does not. In practice the phase != "dispatched" guard already blocks reaching this while handles are detached (checkpointing requires the idle phase), so this is a defense-in-depth consistency nit rather than a live bug. Aligning both keeps the two frontends symmetric if the phase semantics ever change.

♻️ Optional consistency tweak

     def get_combine_payload_tensor_in_workspace(
         self, runtime_max_tokens_per_rank: int, hidden_size: int, dtype: torch.dtype
     ) -> torch.Tensor:
         ...
+        self._require_mapped()
         if self._dispatch_state.get("phase") != "dispatched":
             raise RuntimeError(
                 "get_combine_payload_tensor_in_workspace called before a successful dispatch"
             )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py`
around lines 703 - 722, Update get_combine_payload_tensor_in_workspace to call
the existing _require_mapped() guard before checking the dispatch phase,
matching MoeAlltoAll’s implementation while preserving the current phase
validation and behavior.

tests/unittest/_torch/test_mnnvl_memory_lifecycle.py (1)

290-349: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add direct coverage for the new MnnvlMoe lifecycle statics.

The tests cover MnnvlMemory and the two communication wrappers. They never call the real MnnvlMoe.checkpoint_prepare, MnnvlMoe.checkpoint_restore, or MnnvlMoe.require_mapped. Line 339 replaces require_mapped with a Mock, so its logic stays untested.

Three branches in tensorrt_llm/_mnnvl_utils.py lines 655-685 have no coverage: the None guard for each workspace, the conditional moe_initialize_workspace call when moe_workspace_tensor is set, and the per-workspace mapped check. Monkeypatched workspace mocks are enough, matching the pattern at lines 302-304.

💚 Suggested additional tests

def test_moe_require_mapped_rejects_detached_prepare_workspace(monkeypatch):
    monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", Mock(mapped=True))
    monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", Mock(mapped=False))

    with pytest.raises(RuntimeError, match="workspace handles are unmapped"):
        mnnvl.MnnvlMoe.require_mapped()


def test_moe_checkpoint_prepare_skips_unallocated_workspaces(monkeypatch):
    workspace = Mock(mapped=True)
    monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", workspace)
    monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", None)

    mnnvl.MnnvlMoe.checkpoint_prepare()

    workspace.checkpoint_prepare.assert_called_once_with()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/test_mnnvl_memory_lifecycle.py` around lines 290 - 349,
Add direct tests for MnnvlMoe.checkpoint_prepare, MnnvlMoe.checkpoint_restore,
and MnnvlMoe.require_mapped using monkeypatched workspace objects. Cover None
workspaces, conditional moe_initialize_workspace when moe_workspace_tensor is
set, per-workspace mapped validation, and delegation to checkpoint methods;
remove the require_mapped mock in
test_two_sided_combine_requires_new_prepare_before_next_dispatch so the real
logic is exercised.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes